#!/usr/bin/python

R_1_3 = {
  'A': 'ALA',
  'C': 'CYS',
  'D': 'ASP',
  'E': 'GLU',
  'F': 'PHE',
  'G': 'GLY',
  'H': 'HIS',
  'I': 'ILE',
  'K': 'LYS',
  'L': 'LEU',
  'M': 'MET',
  'N': 'ASN',
  'P': 'PRO',
  'Q': 'GLN',
  'R': 'ARG',
  'S': 'SER',
  'T': 'THR',
  'V': 'VAL',
  'W': 'TRP',
  'Y': 'TYR',
}

R_3_1 = {
  'ALA' : 'A',
  'CYS' : 'C',
  'ASP' : 'D',
  'GLU' : 'E',
  'PHE' : 'F',
  'GLY' : 'G',
  'HIS' : 'H',
  'ILE' : 'I',
  'LYS' : 'K',
  'LEU' : 'L',
  'MET' : 'M',
  'ASN' : 'N',
  'PRO' : 'P',
  'GLN' : 'Q',
  'ARG' : 'R',
  'SER' : 'S',
  'THR' : 'T',
  'VAL' : 'V',
  'TRP' : 'W',
  'TYR' : 'Y',
  'NTR' : 'X',  # NTR and CTR and HOH
  'CTR' : 'Z',
  'HOH' : 'O',
  '_CA' : 'B',
  'RSB' : 'U'
}

# '-':-1, '0':0, '+':+1, 'm':Dummy
POSS_IONI = {'ASP':['-', '0'], 'RSB':['+', '0'], 'ARG':['+', '0'], 'GLU':['-', '0'], 'HOH':['0', '-', '+']}

POSS_CRG = {'ASP':[-1, 0], 'RSB':[1, 0], 'ARG':[1, 0], 'GLU':[-1, 0], 'HOH':[0, -1, 1], 'HYD':[0, 1]}
ION_HEAD  = {'-':'-', '0':'0', '+':'+', 'D':'m'}

SYMBOLTOINT = {'-':-1, '+':1, '0':0}
INTTOSYMBOL = {-1:'-', 1:'+', 0:'0'}

import os, sys

def load_key_res(ifile):
    """ load residue names from a file, such as "ASPA0085" """
    key_res = []
    lines = open(ifile).readlines()
    for line in lines:
        if line.strip()[0] != '#':
            key_res.append(line[:-1])

    return key_res

def com_list(lst1, lst2):
    """combine two lists"""
    # empty list
    if not lst1: return lst2
    if not lst2: return lst1

    all_state = [a+b for a in lst1 for b in lst2]

    return all_state

def set_runs(residues):
    short_res_name = []
    sub_state = []       # all the possible ionization states of the residues
    for res in residues:
        short_name = R_3_1[res[:3]]
     # don't consider the chain id
     #   short_name += res[3]
        short_name += str(int(res[4:]))
        short_res_name.append(short_name)

        sub_state = com_list(sub_state, POSS_IONI[res[:3]])
        
    for sub in sub_state:
        dir_name = ''
        for i in range(len(residues)):
            dir_name += short_res_name[i]
            dir_name += sub[i]
        if not os.path.exists(dir_name):
            os.mkdir(dir_name)
        os.chdir(dir_name) 
        os.system('ln -s ../energies.opp .')
        os.system('ln -s ../run.prm .')
        os.system('ln -s ../extra.tpl .')
        os.system('ln -s ../ms_gold .')

        os.system('cp ../head3.lst .')
        lines = open('head3.lst').readlines()
        new_line = []
        new_line.append(lines[0])
        lines.pop(0)
        for line in lines:
            if line[6:9] not in POSS_IONI.keys():
                 new_line.append(line)
                 continue
             
            buff_name = R_3_1[line[6:9]] + str(int(line[12:16]))
            if (buff_name in short_res_name) and (line[9] != ION_HEAD[sub[short_res_name.index(buff_name)]]):
                new_line.append(line[:21] + 't' + line[22:])
            else:
                new_line.append(line)

        open('head3.lst', 'w').writelines(new_line)

        os.system('cp /home/xzhu/pfile/submit_ms.sh .')
        os.system('qsub submit_ms.sh')

        os.chdir('..')

def retrieve_runs(residues):
    short_res_name = []
    sub_state = []       # all the possible ionization states of the residues

    f_collt = open('ener.txt', 'w')
    head_str = ''
    for res in residues:
        head_str += ('%-10s' % res)
    head_str += '%-9s%-9s%-9s%-9s' % ('Avg', 'std', 'occ', 'Nms')
    f_collt.write(head_str + '\n')   

    for res in residues:
        short_name = R_3_1[res[:3]]
     # don't consider the chain id
     #   short_name += res[3]
        short_name += str(int(res[4:]))
        short_res_name.append(short_name)

        sub_state = com_list(sub_state, POSS_IONI[res[:3]])

    for sub in sub_state:
        dir_name = ''
        for i in range(len(residues)):
            dir_name += short_res_name[i]
            dir_name += sub[i]
        if not os.path.exists(dir_name):
            print "no directory: " + dir_name
            continue
        os.chdir(dir_name)
        

        if not os.path.exists('ms_crg'):
            os.system('mpirun -n 6 ~/bin/e_crg/te -f ../f.txt')
        ms_fields = (open('ms_crg').readlines()[1]).split()
        line_to_write = ''
        n_res = len(residues)
        for i in range(n_res):
            line_to_write += ('%-10s' % INTTOSYMBOL[int(ms_fields[i])])
        line_to_write += ('%-9.1f%-9.1f%-9.1f%-9d' % (float(ms_fields[n_res+3]), float(ms_fields[n_res+4]), \
                          float(ms_fields[n_res+5][:-1]) / 100, int(ms_fields[n_res+6])))
        f_collt.write(line_to_write + '\n')
        
        os.chdir('..') 
        
def main(fl_name):
    all_res = load_key_res(fl_name)
#    set_runs(all_res)
    retrieve_runs(all_res) 

if __name__ == '__main__':
    main(sys.argv[1])
